import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.util.Collector;
import org.apache.flink.util.OutputTag;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;

public class WordCount {
    private static final Logger logger = LoggerFactory.getLogger(WordCount.class);

    static final OutputTag<String> wantedPersonTag = new OutputTag<String>("wantedPerson") {
    };

    public static void main(String[] args) throws Exception {
        // set up the execution environment
        final StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(new Configuration());

        // read data from socket
        DataStream<String> textStream = env.socketTextStream("localhost", 9999, "\n")
                .name("MySocket").uid("MySocket");

        // split up the lines in pairs (2-tuples) containing: (word,1)
        SingleOutputStreamOperator<Tuple2<String, Integer>> tokenStream =
        //      textStream.flatMap(new Tokenizer())
                textStream.process(new Tokenizer2())
                        .name("MyTokenizer2").uid("MyTokenizer2")
                        .setMaxParallelism(1024).setParallelism(2);

        DataStream<Tuple2<String, Integer>> counts =
                tokenStream
                        // group by the tuple field "0" and sum up tuple field "1"
                        .keyBy(value -> value.f0)
                        .sum(1)
                        .name("MySummary").uid("MySummary");

        tokenStream.getSideOutput(wantedPersonTag)
                .print().name("WantedPersonPrint").uid("WantedPersonPrint");

        // output the result to stdout
        counts.print().name("MyPrint").uid("MyPrint");

        env.execute("WordCount");
    }

    public static final class Tokenizer implements FlatMapFunction<String, Tuple2<String, Integer>> {
        @Override
        public void flatMap(String value, Collector<Tuple2<String, Integer>> out) {
            // normalize and split the line
            String[] tokens = value.toLowerCase().split("\\W+");

            // emit the pairs
            for (String token : tokens) {
                if (token.length() > 0) {
                    out.collect(new Tuple2<>(token, 1));
                }
            }
        }
    }

    private static class Tokenizer2 extends ProcessFunction<String, Tuple2<String, Integer>> {
        private transient List<String> wantedPersons;

        @Override
        public void open(Configuration parameters) throws Exception {
            super.open(parameters);

            this.wantedPersons = new ArrayList<>();
            wantedPersons.add("jack");
        }

        @Override
        public void processElement(String value, ProcessFunction<String, Tuple2<String, Integer>>.Context ctx,
                                   Collector<Tuple2<String, Integer>> out) {
            // normalize and split the line
            String[] tokens = value.toLowerCase().split("\\W+");

            // emit the pairs
            for (String token : tokens) {
                if (wantedPersons.contains(token)) {
                    ctx.output(wantedPersonTag, token);
                } else if (token.length() > 0) {
                    out.collect(new Tuple2<>(token, 1));
                }
            }
        }
    }
}
